{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"align-center\">\n",
    "<a href=\"https://oumi.ai/\"><img src=\"https://oumi.ai/docs/en/latest/_static/logo/header_logo.png\" height=\"200\"></a>\n",
    "\n",
    "[![Documentation](https://img.shields.io/badge/Documentation-latest-blue.svg)](https://oumi.ai/docs/en/latest/index.html)\n",
    "[![Discord](https://img.shields.io/discord/1286348126797430814?label=Discord)](https://discord.gg/oumi)\n",
    "[![GitHub Repo stars](https://img.shields.io/github/stars/oumi-ai/oumi)](https://github.com/oumi-ai/oumi)\n",
    "</div>\n",
    "\n",
    "👋 Welcome to Open Universal Machine Intelligence (Oumi)!\n",
    "\n",
    "🚀 Oumi is a fully open-source platform that streamlines the entire lifecycle of foundation models - from [data preparation](https://oumi.ai/docs/en/latest/resources/datasets/datasets.html) and [training](https://oumi.ai/docs/en/latest/user_guides/train/train.html) to [evaluation](https://oumi.ai/docs/en/latest/user_guides/evaluate/evaluate.html) and [deployment](https://oumi.ai/docs/en/latest/user_guides/launch/launch.html). Whether you're developing on a laptop, launching large scale experiments on a cluster, or deploying models in production, Oumi provides the tools and workflows you need.\n",
    "\n",
    "🤝 Make sure to join our [Discord community](https://discord.gg/oumi) to get help, share your experiences, and contribute to the project! If you are interested in joining one of the community's open-science efforts, check out our [open collaboration](https://oumi.ai/community) page.\n",
    "\n",
    "⭐ If you like Oumi and you would like to support it, please give it a star on [GitHub](https://github.com/oumi-ai/oumi)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Distillation Overview\n",
    "\n",
    "In this tutorial, we'll show how we trained MiniMath-R1-1.5B!\n",
    "\n",
    "We'll use the Oumi framework to streamline the process and achieve high-quality results.\n",
    "\n",
    "We'll cover the following topics:\n",
    "1. Prerequisites\n",
    "2. Model and Data Preparation\n",
    "3. Fine-Tuning\n",
    "4. Evaluation\n",
    "5. Upload to HuggingFace"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Prerequisites\n",
    "\n",
    "## Hardware\n",
    "The defaults in this tutorial are scaled down for demonstration purposes.\n",
    "\n",
    "The true values are left to code comments within each section.\n",
    "\n",
    "We recommend 8xA100-80GB GPUs to complete in a timely manner with adequate performance.\n",
    "\n",
    "## Oumi Installation\n",
    "\n",
    "First, let's install Oumi and vLLM. You can find more detailed instructions [here](https://oumi.ai/docs/en/latest/get_started/installation.html). Here, we include Oumi's GPU dependencies.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install oumi[gpu]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Creating our working directory\n",
    "For our experiments, we'll use the following folder to save the model, training artifacts, and our working configs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "\n",
    "tutorial_dir = \"distillation_tutorial\"\n",
    "\n",
    "Path(tutorial_dir).mkdir(parents=True, exist_ok=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup the environment\n",
    "\n",
    "We'll need to set the following environment variables:\n",
    "- [Optional] HF_TOKEN: Your [HuggingFace](https://huggingface.co/docs/hub/en/security-tokens) token, in case you want to access a private model.\n",
    "- [Optional] WANDB_API_KEY: Your [wandb](https://wandb.ai) token, in case you want to log your experiments to wandb."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "os.environ[\"HF_TOKEN\"] = \"INSERT TOKEN HERE\"\n",
    "os.environ[\"WANDB_API_KEY\"] = \"INSERT API KEY HERE\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Model and Data Preparation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Model Download\n",
    "\n",
    "For our purposes it will be much faster if we download our models first.\n",
    "\n",
    "We'll use the `hf_transfer` package to download."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install hf_transfer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!HF_HUB_ENABLE_HF_TRANSFER=1 hf download \\\n",
    "    deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B --exclude original/*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Baseline Evals\n",
    "\n",
    "Before we can improve our small model, we should measure a baseline.\n",
    "\n",
    "The below code will run the MMLU PRO Math task from LM Harness. \n",
    "\n",
    "Note that this will take some time, so we've recorded our results below for your convenience:\n",
    "\n",
    "| Model | MMLU Pro Math Accuracy |\n",
    "|-------|------------------------|\n",
    "| R1 Distill 1.5B | 38.49% +- 1.32% |"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Run Evals"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile $tutorial_dir/eval_small.yaml\n",
    "\n",
    "model:\n",
    "  model_name: \"deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B\"\n",
    "  torch_dtype_str: \"bfloat16\"\n",
    "  # shard_for_eval: True # Uncomment this line for multi-gpu setups.\n",
    "\n",
    "\n",
    "tasks:\n",
    "  - evaluation_backend: lm_harness\n",
    "    task_name: mmlu_pro_math\n",
    "\n",
    "output_dir: \"distillation_tutorial/output/evaluation\"\n",
    "generation:\n",
    "  batch_size: 1 # LM Harness recommends BS=1 for reproducibility.\n",
    "  # batch_size: 128  # Replace with 256 for 8xA100-80GB"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!oumi evaluate -c \"$tutorial_dir/eval_small.yaml\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Prepare Training Data\n",
    "\n",
    "Oumi has released an Apache 2.0 license math dataset at `oumi-ai/MetaMathQA-R1`, let's go ahead and download it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!HF_HUB_ENABLE_HF_TRANSFER=1 hf download oumi-ai/MetaMathQA-R1 \\\n",
    "    --exclude original/* --repo-type dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Fine-Tuning\n",
    "\n",
    "Now that the data is downloaded, we can begin fine-tuning the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile $tutorial_dir/train.yaml\n",
    "\n",
    "model:\n",
    "  model_name: \"deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B\"\n",
    "  trust_remote_code: true\n",
    "  torch_dtype_str: \"bfloat16\"\n",
    "  model_max_length: 4096\n",
    "  device_map: \"auto\"\n",
    "\n",
    "data:\n",
    "  train:\n",
    "    datasets:\n",
    "      - dataset_name: \"PromptResponseDataset\"\n",
    "        split: \"train\"\n",
    "        sample_count: 25000  # 25k samples is enough to get the desired effect\n",
    "        dataset_kwargs: {\n",
    "          \"hf_dataset_path\": \"oumi-ai/MetaMathQA-R1\",\n",
    "          \"prompt_column\": \"prompt\",\n",
    "          \"response_column\": \"response\",\n",
    "        }\n",
    "        shuffle: True\n",
    "        seed: 42\n",
    "    seed: 42\n",
    "\n",
    "training:\n",
    "  output_dir: \"distillation_tutorial/output/finetune\"\n",
    "\n",
    "  # For a single GPU, the following gives us a batch size of 8\n",
    "  # If training with multiple GPUs, feel free to reduce gradient_accumulation_steps\n",
    "  per_device_train_batch_size: 1\n",
    "  gradient_accumulation_steps: 8\n",
    "  \n",
    "  # ***NOTE***\n",
    "  # We set it to 10 steps to first verify that it works\n",
    "  # Comment out the line below to have it train for 1 full epoch (all the data) instead.\n",
    "  # Note: 1 full epoch will take about 13 minutes on 8xH100-80GB.\n",
    "  max_steps: 10\n",
    "  num_train_epochs: 1\n",
    "  learning_rate: 1e-4\n",
    "  warmup_ratio: 0.1\n",
    "  logging_steps: 10\n",
    "  save_steps: 0\n",
    "  max_grad_norm: 10\n",
    "  weight_decay: 0.01\n",
    "  compile: False\n",
    "\n",
    "  \n",
    "  trainer_type: \"TRL_SFT\"\n",
    "  optimizer: \"adamw_torch_fused\"\n",
    "  enable_gradient_checkpointing: True\n",
    "  gradient_checkpointing_kwargs:\n",
    "    use_reentrant: False\n",
    "  ddp_find_unused_parameters: False\n",
    "  dataloader_num_workers: \"auto\"\n",
    "  dataloader_prefetch_factor: 32\n",
    "  empty_device_cache_steps: 1\n",
    "\n",
    "# Uncomment this for distributed training\n",
    "# fsdp:\n",
    "#   enable_fsdp: True\n",
    "#   backward_prefetch: \"BACKWARD_POST\"\n",
    "#   forward_prefetch: True\n",
    "#   cpu_offload: True\n",
    "#   auto_wrap_policy: \"TRANSFORMER_BASED_WRAP\"\n",
    "#   transformer_layer_cls: \"Qwen2DecoderLayer\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Single GPU"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!oumi train -c \"$tutorial_dir/train.yaml\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Multi-GPU"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!oumi distributed torchrun -m oumi train -c \"$tutorial_dir/train.yaml\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Evaluation\n",
    "\n",
    "Now that we have a new distilled model, let's evaluate it on the same benchmark."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile $tutorial_dir/eval_small_fft.yaml\n",
    "\n",
    "model:\n",
    "  model_name: \"./distillation_tutorial/output/\"\n",
    "  torch_dtype_str: \"bfloat16\"\n",
    "  # shard_for_eval: True # Uncomment this line for multi-gpu setups.\n",
    "\n",
    "\n",
    "tasks:\n",
    "  - evaluation_backend: lm_harness\n",
    "    task_name: mmlu_pro_math\n",
    "\n",
    "output_dir: \"distillation_tutorial/output/evaluation\"\n",
    "generation:\n",
    "  batch_size: 1 # LM Harness recommends BS=1 for reproducibility.\n",
    "  # batch_size: 256  # Replace with 256 for 8xA100-80GB"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!oumi evaluate -c \"$tutorial_dir/eval_small_fft.yaml\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Results\n",
    "\n",
    "After we finetuned the model following the steps above, we achieved the following results:\n",
    "\n",
    "| Model           | Accuracy        |\n",
    "|-----------------|-----------------|\n",
    "| R1 Distill 1.5B | 38.49% +- 1.32% |\n",
    "| MiniMath R1 1.5B | 44.4% +- 1.34% |"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Upload to HuggingFace\n",
    "\n",
    "After fine-tuning, let's upload our model to HuggingFace to make it easily portable to other places."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "HUGGINGFACE_REPO_PATH = \"your-user-name/your-model-name\"\n",
    "LOCAL_MODEL_PATH = f\"./{tutorial_dir}/output\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Upload Model\n",
    "\n",
    "Transformers makes it fairly easy to upload the model itself."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import transformers\n",
    "\n",
    "model = transformers.AutoModel.from_pretrained(LOCAL_MODEL_PATH, torch_dtype=\"bfloat16\")\n",
    "model.push_to_hub(HUGGINGFACE_REPO_PATH)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Upload Configs\n",
    "\n",
    "HuggingFace by default doesn't upload a number of important configs for inference so we\n",
    "have to upload these manually."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from huggingface_hub import HfApi\n",
    "\n",
    "api = HfApi()\n",
    "model_files = [f for f in Path(LOCAL_MODEL_PATH).glob(\"*.json\")]\n",
    "\n",
    "for file in model_files:\n",
    "    file_name = file.name\n",
    "    api.upload_file(\n",
    "        path_or_fileobj=file,\n",
    "        path_in_repo=file_name,\n",
    "        repo_id=HUGGINGFACE_REPO_PATH,\n",
    "        repo_type=\"model\",\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 🧭 What's Next?\n",
    "\n",
    "Congrats on finishing this notebook! Feel free to check out our other [notebooks](https://github.com/oumi-ai/oumi/tree/main/notebooks) in the [Oumi GitHub](https://github.com/oumi-ai/oumi), and give us a star! You can also join the Oumi community over on [Discord](https://discord.gg/oumi).\n",
    "\n",
    "📰 Want to keep up with news from Oumi? Subscribe to our [Substack](https://blog.oumi.ai/) and [Youtube](https://www.youtube.com/@Oumi_AI)!\n",
    "\n",
    "⚡ Interested in building custom AI in hours, not months? Apply to get [early access](https://oumi-ai.typeform.com/early-access) to the Oumi Platform, or [chat with us](https://calendly.com/d/ctcx-nps-47m/chat-with-us-get-early-access-to-the-oumi-platform) to learn more!"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "oumi",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
